const char *help = "\
progname: fft2D.cc\n\
code2html: This program reads a pgm image and computes its FFT 2D.\n\
version: Torch3 vision2.0, 2005\n\
(c) Sebastien Marcel (marcel@idiap.ch)\n";


#include "ImageGray.h"
#include "ipFFT2D.h"
#include "tiffDiskXFile.h"
#include "ImageRgb.h"
#include "DiskXFile.h"
#include "CmdLine.h"
#include "fft2d_util.h"
#include "xtprobeImageDiskXFile.h"

using namespace Torch;

//******   Pré-traitement de l'image *******//

ImageGray * treatImage(Image *source,char* name)
{	

	Image *image_in = new ImageGray();
	image_in->copyFrom(source);

	// check if the image if power of 2
	int new_width = (int) pow(2.0, ceil(log2((double)image_in->width)));
	int new_height = (int) pow(2.0, ceil(log2((double)image_in->height)));

	print("width = %d -> new width = %d\n", image_in->width, new_width);
	print("height = %d -> new height = %d\n", image_in->height, new_height);

	// re-copy the input image into an image of size power of 2
	ImageGray *newimage = new ImageGray(new_width, new_height);
	pasteGray(image_in->data, 0, 0, image_in->width, image_in->height, newimage->data, new_width, new_height);
	newimage->updatePixmapFromData();
	newimage->save(name);

     return newimage;
}

//******   Construction du nom de fichier *******//

char * setTitle(char* string,char* prefix)
{
	
	char *str = new char[300];
	char *str1 = new char[300];
	int i, j=0;
	
// Extraction du nom du fichier sans extension

	for (i = 0; string[i] != '\0'; i ++)
	 {
           if(string[i]!='/')
            {
  	     str[j] = string[i];
             j++;
            }
	   else
            {
	    j=0;
             }             
         }
	str[j ++] = (char)NULL;

	for (i = 0; str[i] != '.'; i ++)
	 {
          str1[i] = str[i];
         }
            
        str1[i] = (char)NULL;	 
 
	strcat(str1, prefix);
	return str1;

}

//******   Transformation de fourier normale *******//

void normalFFT(ipCore *fft2d,Image *image,char* reelle,char* imaginaire,bool verbose,bool swap,bool improve,bool highpass,bool lowpass)
{	
	//** FFT2D   Transformation Normale **//
     
	print("Computing FFT 2D ...\n");
	fft2d->process(image);

	// print coefficients
	real *fft_r = fft2d->seq_out->frames[0];
	real *fft_i = fft2d->seq_out->frames[1];

	real energy = 0;
	int n_hpass = 0;
	int n_lpass = 0;
	if(verbose) print("FFT 2D:\n");	
	for(int i = 0; i < fft2d->seq_out->frame_size; i++)
	{
		if(verbose) print("[%d] = %g (%g)\n", i, fft_r[i], fft_i[i]);

		energy += fft_r[i] * fft_r[i];

		if(i > 0)
		{
			if(fft_r[i] > highpass) n_hpass++;
			if(fft_r[i] < lowpass) n_lpass++;
		}
	}

	energy = sqrt(energy);
	real log_energy = energy < 1.0 ? 0.0 : (real) log(energy); 

	print("DC frequency = %g + %gi\n", fft_r[0], fft_i[0]);
	print("Nyquist frequency = %g + %gi\n", fft_r[fft2d->seq_out->frame_size/2], fft_i[fft2d->seq_out->frame_size/2]);
	print("log Energy = %g\n", log_energy);
	print("number of high passed = %d\n", n_hpass);
	print("number of low passed = %d\n", n_lpass);


	// Conversion des tableaux 1D en images 2D

	Image *resultImg = new ImageGray();
	resultImg = Tab1D_to_Image2D (fft_r, image->width, image->height, improve, swap);
	resultImg->save(reelle);

	resultImg = Tab1D_to_Image2D (fft_i, image->width, image->height, improve, swap);
	resultImg->save(imaginaire);
        
	delete resultImg;

}


//******   Transformation de fourier inverse *******//

void inverseFFT(ipCore *fft2d,Image *image,Image *source,char* invers,bool verbose,bool clean)
{	

	// IFFT2D  Transformation Inverse

  		ipCore *ifft2d = NULL;
	
		ifft2d = new ipFFT2D(image->width, image->height, true);
		//ifft2d->setBOption("verbose", verbose);
	
		print("Computing IFFT 2D ...\n");
		ifft2d->process(fft2d->seq_out);

		// compute reconstruction error
		real rmse = 0.0;
		
		if(verbose) print("IFFT 2D:\n");	
		for(int i = 0; i < ifft2d->seq_out->frame_size ; i++)
		{
			if(verbose) print("[%d] = %g (%g)\n", i, ifft2d->seq_out->frames[0][i], ifft2d->seq_out->frames[1][i]);
		   	real z = image->frames[0][i] - ifft2d->seq_out->frames[0][i];
			rmse += z * z;
		}
		print("RMSE = %g\n", rmse);


	//** Enregistrement de l'image **//

		ImageGray *image_out = new ImageGray(source->width, source->height);
		cropGray(ifft2d->seq_out->frames[0], image->width, image->height, 0,0, source->width, source->height, image_out->data);
		image_out->updatePixmapFromData();		
		image_out->save(invers);


		delete ifft2d;
		
}

//******   Application des filtres aux images *******//

ipCore * Filtrage(ipCore *fft2d,Image *image,bool swap,bool improve,int filtre,bool inverse2,real facteur,char* name11,char* name22)
{
	real *fft_r = fft2d->seq_out->frames[0];
	real *fft_i = fft2d->seq_out->frames[1];

	real *filt;
	if (filtre == 1)
	{
		filt = DefinirFiltre1 (image->width, image->height, facteur);
		AppliquerFiltre1 (fft_r, filt, image->width*image->height);
		AppliquerFiltre1 (fft_i, filt, image->width*image->height);
	}
	
	if (filtre == 2)
	{
		filt = DefinirFiltre2 (image->width, image->height, facteur, inverse2);
		AppliquerFiltre2 (fft_r, filt, image->width*image->height);
		AppliquerFiltre2 (fft_i, filt, image->width*image->height);
	}
	
	Image *resultImg = new ImageGray();
	resultImg = Tab1D_to_Image2D (fft_i, image->width, image->height, improve, swap);
	resultImg->save(name22);
	resultImg = Tab1D_to_Image2D (fft_r, image->width, image->height, improve, swap);
	resultImg->save(name11);

return fft2d;
}

//******   Retrait du filtre 1 sur l'image *******//

ipCore * backFiltrage(ipCore *fft2d,Image *image,bool swap,bool improve,real facteur,char* name11,char* name22)
{
	real *fft_r = fft2d->seq_out->frames[0];
	real *fft_i = fft2d->seq_out->frames[1];

	real *filt;

	filt = DefinirFiltre1 (image->width, image->height, facteur);
	int *posr=AppliquerBackFiltre1 (fft_r, filt, image->width*image->height);
	int *posi=AppliquerBackFiltre1 (fft_i, filt, image->width*image->height);

	fft_i = noiseRemove(fft_i, image->width, image->height, improve, swap,posi);
	fft_r = noiseRemove(fft_r, image->width, image->height, improve, swap,posr);

	Image *resultImg = new ImageGray();
	resultImg = Tab1D_to_Image2D(fft_i, image->width, image->height, improve, swap);
	resultImg->save(name22);
	resultImg = Tab1D_to_Image2D(fft_r, image->width, image->height, improve, swap);
	resultImg->save(name11); 

return fft2d;
}



//******   Traitement de l'image *******//

void processImage(Image *source,char *image_filename,char *finalname,bool swap,bool improve,int filtre,bool inverse,real facteur,bool verbose,real highpass,real lowpass,bool clean)
{	


//******  Nommage des images  ***********//

	char *suffix = new char[300]; 
	char *aux = new char[300];
     	strcpy(suffix, "_");
        if(clean) strcpy(suffix, "_BACK_");

	char *name = new char[300];  
	strcpy(aux, suffix);
	strcat(aux, "TREAT.pgm"); 
	strcpy(name, setTitle(image_filename,aux));
	ImageGray *newimage=treatImage(source,name);

	char *name1 = new char[300];  
	strcpy(aux, suffix);
	strcat(aux, "FFTR.pgm"); 
	strcpy(name1, setTitle(image_filename,aux));

	char *name2 = new char[300];  
	strcpy(aux, suffix);
	strcat(aux, "FFTI.pgm"); 
	strcpy(name2, setTitle(image_filename,aux));

	char *name11 = new char[300];  
	strcpy(aux, suffix);
	strcat(aux, "Filtre_FFTR.pgm"); 
	strcpy(name11, setTitle(image_filename,aux));

	char *name22 = new char[300];  
	strcpy(aux, suffix);
	strcat(aux, "Filtre_FFTI.pgm"); 
	strcpy(name22, setTitle(image_filename,aux));


//******  Transformation de Fourier  ***********//


	ipCore *fft2d = new ipFFT2D(newimage->width, newimage->height);
	fft2d->setBOption("verbose", verbose);

       	normalFFT(fft2d,newimage,name1,name2,verbose,swap,improve,highpass,lowpass);
	
	ipCore *fft2di = new ipFFT2D(newimage->width, newimage->height);

        if((!clean)&&(filtre!=0))
        {
        // Filtrage normale
	fft2di = Filtrage(fft2d,newimage,swap,improve,filtre,inverse,facteur,name11,name22);
	inverseFFT(fft2d,newimage,source,finalname,verbose,clean); 

	}
	else if(clean)
	{
        // Revenir à l'image de départ : Retrait du filtrage

	fft2di = backFiltrage(fft2d,newimage,swap,improve,facteur,name11,name22);
	inverseFFT(fft2d,newimage,source,finalname,verbose,clean); 
	}

	
 
}

////////////************* PROGRAMME PRINCIPALE ******************///////////////////////

int main(int argc, char **argv)
{
	char *image_filename;
	bool verbose;
	bool inverse;
	bool swap;
	bool improve;
	bool clean;
	real highpass;
	real lowpass;
	real facteur;
	int filtre;
	

  	xtprobeImageDiskXFile *image_file = NULL;
	Image *source = NULL;


	CmdLine cmd;
	cmd.setBOption("write log", false);
  	cmd.info(help);
  	cmd.addText("\nArguments:");
  	cmd.addSCmdArg("image filename", &image_filename, "image filename");
  	cmd.addText("\nOptions:");
 	cmd.addBCmdOption("-verbose", &verbose, false, "verbose");
  	cmd.addBCmdOption("-inverse", &inverse, false, "inverse");
  	cmd.addRCmdOption("-hpass", &highpass, 100.0, "high pass");
  	cmd.addRCmdOption("-lpass", &lowpass, -100.0, "low pass");
  	cmd.addBCmdOption("-improve", &improve, false, "improve");
  	cmd.addBCmdOption("-swap", &swap, false, "swap");
 	cmd.addBCmdOption("-clean", &clean, false, "Nettoyer le filtre 1");
  	cmd.addICmdOption("-filtre", &filtre, 0, "filtre");
  	cmd.addRCmdOption("-facteur", &facteur, 0.5, "facteur");
	cmd.read(argc, argv);


// Création et chargement de l'image

	source = new ImageRgb();  	
	source->setBOption("verbose", verbose);

    // Choix: xtprobeImageDiskXFile car il permet de travailler avec plusieurs types d'images 

	image_file = new xtprobeImageDiskXFile(image_filename, "r");
	source->loadImageXFile(image_file);
	delete image_file;


	if(verbose)
	{
		print("Image info:\n");
		print("   width = %d\n", source->width);
		print("   height = %d\n", source->height);
		print("   format = %s (%d)\n", source->coding, source->n_planes);
	}

    
//******  Transformation de Fourier  ***********//

	char *name = new char[300];
 	if(!clean)
        	strcpy(name, setTitle(image_filename,(char*)"_MODIF.pgm"));
	else	strcpy(name, setTitle(image_filename,(char*)"_CLEAN.pgm"));

	processImage(source,image_filename,name,swap,improve,filtre,inverse,facteur,verbose,highpass,lowpass,clean);
	
	delete source;

	return(0);
}
